/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "OpenSparseBLAS_func.h"
#include <complex>
#include <vector>


/*****************csc spsv****************/

namespace _spblasinl {
    template <typename T>
    inline T abs_val(const T& v) { return std::fabs(v); }
    template <typename T>
    inline T abs_val(const std::complex<T>& v) { return std::abs(v); }

    /* 安全下标 (仅 debug 用，可改 assert) */
    inline bool in_range(int idx, int len) { return idx >= 0 && idx < len; }

    /* 指针合法性简单检查 */
    template <typename... Ptrs>
    inline bool all_nonnull(Ptrs... ps) {
        return (... && (ps != nullptr));
    }
} 

// A * x = b
template <typename T>
inline OpenSpB_status OpenSpB_spsv(int m,
                            const int *csc_col_ptr, 
                            const int *csc_row_ind,
                            const T *csc_vals,
                            T *x,
                            const T *b,
                            OpenSpB_base_type index_base,
                            OpenSpB_uplo_type uplo)
{
    using namespace _spblasinl;

    if (m == 0)               return OPENSPBLAS_STATUS_SUCCESS;
    if (!all_nonnull(csc_col_ptr, csc_row_ind, csc_vals, b, x))
        return OPENSPBLAS_STATUS_INVALID;

    int base = (index_base == OPENSPBLAS_ZERO_BASE) ? 0 : 1;

    const int bsz = static_cast<int>(base);   // 0 or 1

    /* 先把 b 拷贝到工作向量 work，之后再写回 x */
    std::vector<T> work(b, b + m);

    /*================== 前向：LOWER  ====================================*/
    if (uplo == OPENSPBLAS_LOWER)
    {
        for (int j = 0; j < static_cast<int>(m); ++j)
        {
            T diag{};                  // 0 值
            const int col_start = csc_col_ptr[j]   - bsz;
            const int col_end   = csc_col_ptr[j+1] - bsz;

            /* 第 1 遍扫描：扣减已知列贡献，并找到对角 */
            for (int p = col_start; p < col_end; ++p) {
                const int row = csc_row_ind[p] - bsz;
                if (row == j) diag = csc_vals[p];
                else if (row  < j) {
                    work[j] -= csc_vals[p] * x[row];
                }
            }
            /* 对角合法？ */
            if (abs_val(diag) == T{}) return OPENSPBLAS_STATUS_SINGULAR;

            x[j] = work[j] / diag;

            /* 第 2 遍扫描：更新列下方 rhs */
            for (int p = col_start; p < col_end; ++p) {
                const int row = csc_row_ind[p] - bsz;
                if (row > j) {
                    work[row] -= csc_vals[p] * x[j];
                }
            }
        }
        return OPENSPBLAS_STATUS_SUCCESS;
    }

    /*================== 后向：UPPER  ====================================*/
    if (uplo == OPENSPBLAS_UPPER)
    {
        for (int j = static_cast<int>(m) - 1; j >= 0; --j)
        {
            T diag{};
            const int col_start = csc_col_ptr[j]   - bsz;
            const int col_end   = csc_col_ptr[j+1] - bsz;

            /* 第 1 遍扫描：扣减已知列贡献，并找到对角 */
            for (int p = col_start; p < col_end; ++p) {
                const int row = csc_row_ind[p] - bsz;
                if (row == j) diag = csc_vals[p];
                else if (row  > j) {            // 上三角严格部分
                    work[j] -= csc_vals[p] * x[row];
                }
            }
            if (abs_val(diag) == T{}) return OPENSPBLAS_STATUS_SINGULAR;

            x[j] = work[j] / diag;

            /* 第 2 遍扫描：更新列上方 rhs */
            for (int p = col_start; p < col_end; ++p) {
                const int row = csc_row_ind[p] - bsz;
                if (row < j) {
                    work[row] -= csc_vals[p] * x[j];
                }
            }
        }
        return OPENSPBLAS_STATUS_SUCCESS;
    }
    /* 未知 tri 类型 */
    return OPENSPBLAS_STATUS_INVALID;
}

OpenSpB_status OpenSpB_ssptrsv(int m,
        const int *col_ptr, const int *row_ind, const float *vals,
        float *x, const float *b,
        OpenSpB_base_type base, OpenSpB_uplo_type uplo)
{
    return OpenSpB_spsv<float>(m, col_ptr, row_ind, vals, x, b, base, uplo);
}

OpenSpB_status OpenSpB_dsptrsv(int m,
        const int *col_ptr, const int *row_ind, const double *vals,
        double *x, const double *b,
        OpenSpB_base_type base, OpenSpB_uplo_type uplo)
{
    return OpenSpB_spsv<double>(m, col_ptr, row_ind, vals, x, b, base, uplo);
}

OpenSpB_status OpenSpB_csptrsv(int m,
        const int *col_ptr, const int *row_ind, const void *vals,
        void *x, const void *b,
        OpenSpB_base_type base, OpenSpB_uplo_type uplo)
{
    const auto *v = static_cast<const std::complex<float>*>(vals);
    const auto *rhs = static_cast<const std::complex<float>*>(b);
    auto *sol = static_cast<std::complex<float>*>(x);

    return OpenSpB_spsv<std::complex<float>>(m, col_ptr, row_ind, v, sol, rhs, base, uplo);
}

OpenSpB_status OpenSpB_zsptrsv(int m,
        const int *col_ptr, const int *row_ind, const void *vals,
        void *x, const void *b,
        OpenSpB_base_type base, OpenSpB_uplo_type uplo)
{
    const auto *v = static_cast<const std::complex<double>*>(vals);
    const auto *rhs = static_cast<const std::complex<double>*>(b);
    auto *sol = static_cast<std::complex<double>*>(x);

    return OpenSpB_spsv<std::complex<double>>(m, col_ptr, row_ind, v, sol, rhs, base, uplo);
}